
import os
import shutil
import numpy as np
import random

"""
用于将sum_dataset中的数据切割成
train,test,query,gallery四个部分
训练集:train
验证集:test
测试集:query,gallery
"""

def get_random_file_path(sub_dir_path:str):
    dirlis = os.listdir(sub_dir_path)
    lenth = len(dirlis)
    random_test_idx = int(random.random()*lenth)
    random_test_idx = min(train_size,max(0,random_test_idx))
    filename = dirlis[random_test_idx]
    # 返回文件名称,完整路径名称
    return filename,os.path.join(sub_dir_path,filename)

def path_mkdir_if_notexists(path):
    if os.path.exists(path) == False:
        os.mkdir(path)

def copy_file(src_file:str,dst_file:str):
    try:
        shutil.copyfile(src_file,dst_file)
    except Exception as e:
        print(f"copy file error while : [{src_file}]->[{dst_file}] for {str(e)}")
        exit(0)

def move_dir(src_dir:str,dst_dir:str):
    try:
        shutil.move(src_dir,dst_dir)
    except Exception as e:
        print(f"move dir error while : [{src_dir}]->[{dst_dir}] for {str(e)}")
        exit(0)

def move_file(src_file:str,dst_file):
    try:
        shutil.move(src_file,dst_file)
    except Exception as e:
        print(f"move file error while : [{src_file}]->[{dst_file}] for {str(e)}")
        exit(0)

def part_list_data(part_list:list,ratio:tuple=(8,2)):
    """
    请注意,随机采样分割中part_list不能存在重复的样本
    """
    lenth = len(part_list)
    left_ratio = ratio[0] / sum(ratio)
    left_size = int( lenth * left_ratio )
    right_size = lenth - left_size
    left_list = np.random.choice(part_list,left_size,replace=False).tolist()
    right_list = list( set(part_list) - set(left_list) )
    return left_list,left_size,right_list,right_size


# ----------------------

# 是否跳过phone相关的文件夹,不训练phone类
skip_phone_dir:bool = True

# 获取sum_dataset中的数据组
sub_dir_list = []
lis = os.listdir("sum_dataset")
for one in lis:
    if "phone" in str(one):
        continue
    path = os.path.join("sum_dataset",one)
    if os.path.isdir(path):
        sub_dir_list.append(one)
sub_dir_list.sort()

# 创建data目录
try:
    os.remove("data")
except Exception as e:
    if "系统找不到指定的文件" in str(e):
        # 用于适应win
        pass
    elif "No such file or directory" in str(e):
        # 用于适应linux
        pass
    else:
        print(str(e))
        exit(0)
path_mkdir_if_notexists("data")




# 切割比例:train+test:eval的比例
ratio = ( 8 , 2 )
# gallery:query
eval_ratio = ( 9 ,1 )

# 1
train_list,train_size,eval_list,eval_size = part_list_data(sub_dir_list,ratio)

# train与test文件夹
path_mkdir_if_notexists("data/train")
path_mkdir_if_notexists("data/test")
for sub_dir in train_list:
    src_path = os.path.join("sum_dataset",sub_dir)
    # test:只移动单个文件
    test_filename,test_src_file_path = get_random_file_path(src_path)
    test_dst_dir_path = os.path.join("data/test",sub_dir,)
    path_mkdir_if_notexists(test_dst_dir_path)
    test_dst_file_path = os.path.join(test_dst_dir_path,test_filename)
    move_file(test_src_file_path,test_dst_file_path)
    # train:移动文件夹
    train_dst_path = os.path.join("data/train",sub_dir)
    move_dir(src_path,train_dst_path)
# 布置gallery和query
path_mkdir_if_notexists("data/gallery")
path_mkdir_if_notexists("data/query")
for sub_dir in eval_list:
    src_rot_path = os.path.join("sum_dataset",sub_dir)
    dst_gallery_rot_path = os.path.join("data/gallery",sub_dir)
    path_mkdir_if_notexists(dst_gallery_rot_path)
    dst_query_rot_path = os.path.join("data/query",sub_dir)
    path_mkdir_if_notexists(dst_query_rot_path)
    # 获取src_rot_path中的文件列表
    filelist = os.listdir(src_rot_path)
    glis,gsize,qlis,qsize = part_list_data(filelist,eval_ratio)
    # 移动gallery
    for filename in glis:
        src_path = os.path.join(src_rot_path,filename)
        dst_path = os.path.join(dst_gallery_rot_path,filename)
        move_file(src_path,dst_path)
    # 移动query
    for filename in qlis:
        src_path = os.path.join(src_rot_path,filename)
        dst_path = os.path.join(dst_query_rot_path,filename)
        move_file(src_path,dst_path)
